import copy
import torch
from utils.strategy import logits2span_list
from utils import *
from exp01.comodel import ComModel
from tqdm import tqdm


def test(
        args,
        model: ComModel,
        test_data_loader,
        three_goden_set,
        logger,
        is_test: bool = False,  # 如果是 train 和 dev 数据集，那么 is_test 就是 false
        model_dir='',
):
    logger.info("Test is called! Test begin!")
    model.eval()
    if is_test:  # 如果使用的是 test 数据集
        model.load_state_dict(torch.load(model_dir))
        pass
    aspect_set, opinion_set, triplets_set, pair_set, as_set = set(), set(), set(), set(), set()
    # aspect_set 是一个 set 其中元素是形如 batch_i-start_i-end_i 的字符串
    # pair_set 是 ["i-al-ar-ol-or"]
    # triplets_set 是 [“i-al-ar-ol-or-s"]
    # as_set 是 [”i-al-ar-s“]
    multi_set, single_set = set(), set()
    # multi_set 和 single_set 都是 [“i-al-ar-ol-or-s"]

    multi_aspect_id = three_goden_set[2][3]  # original_data.multi_aspect_id 记录的是所有多标签的 id，从 0 开始

    n_batch_called = 0
    with torch.no_grad():
        for batch_i, test_batch_data in tqdm(enumerate(test_data_loader), desc="testing" if is_test else "validating", delay=0.5):
            n_batch_called += 1
            is_on, count = True, 0  # count 代表的是这一个 test 的样本中执行了多少轮
            inputs_plus = test_batch_data['inputs_plus_for_test']
            token2id = test_batch_data["test_instances"][0].meta_instance.token2id
            token_seq_len = inputs_plus['input_ids'].size(1)

            se_r = []
            inputs_plus_clone = copy.deepcopy(inputs_plus)

            # 开始执行多轮的推理，每次识别一个 aspect 并将结果 (aspect_start_i, aspect_end_i) 放到 se_r 当中
            while is_on:
                model_res = model.get_aspect(inputs_plus_clone)
                as_p, ae_p, _, _, _ = model_res
                # as_p 和 ae_p 的 shape=(batch_size, sentence_seq_len)
                aspect_span = logits2span_list(
                    as_p.squeeze(),
                    ae_p.squeeze(),
                    3,
                    1,
                    args.a,
                    args.a_ww,
                    args.a,
                    token2id,
                    is_test
                )  # a_spans 是一个列表，里边是（aspect_start_i, aspect_end_i）的元组

                if len(aspect_span) == 0:  # 如果炸不到任何一个 aspect 了
                    is_on = False
                    pass
                else:  # 此时应该继续压榨：更新 attention_mask，进一步运算
                    # cal is_on
                    as_input = torch.tensor([0] * (token_seq_len - 1)).unsqueeze(0).to(args.device)
                    as_input[0][aspect_span[0][0]] = 1  # 将这个 aspect_i 的位置标上 1
                    ae_input = torch.tensor([0] * (token_seq_len - 1)).unsqueeze(0).to(args.device)
                    ae_input[0][aspect_span[0][1]] = 1

                    y = model(inputs_plus_clone, as_input, ae_input, args, model_res)  # 但是其实在模型中这个 as_input 和 ae_input 压根就没有被用到
                    # as_input 和 ae_input 压根就没有被用到，真正有用的是上边的 model_res，在 model.forward() 中作为 plus 输入的
                    is_on_logits = y['is_on_logits']  # (batch_size, 2)

                    if count != 0:  # 执行下面的代码代表着 len(a_spans) > 0 且已经检测到了至少一个 aspect 了
                        is_on = torch.argmax(is_on_logits, dim=1).squeeze().item() == 1
                        # is_on 本来应该是(batch_size,) 但是在 test 时，batch_size 为 1，所以 is_on 在这里是一个数字
                        # 又因为 is_on 调用的是 argmax() 方法，这里的含义其实就是说如果 is_on_logits[1] > is_on_logits[0]
                        if not is_on:
                            break
                            pass
                        pass

                    # 这里是确定，使用哪一种方式来更新 se_r 数组，默认 args.use_c=1
                    if args.use_c == 1:  # 默认走这里，
                        se_r = f_join((aspect_span[0][0], aspect_span[0][1]), se_r)
                        # a_spans 是 (aspect_start_i, aspect_end_i) 的列表，但是这个列表目前只有一个元素，因为之前设置的系数是 1
                        # 这里的 use_c 应该是作者进行的优化，通过优化，se_r 的更新方法发生了变化
                        pass
                    else:
                        se_r.append((aspect_span[0][0], aspect_span[0][1]))
                        pass

                    # 现在是
                    inputs_plus_clone['attention_mask'][0][1 + aspect_span[0][0]:1 + aspect_span[0][1] + 1] = 0
                    count += 1

                    if count > 10:  # 如果一个例子已经提取出 10 个信息了，那就直接退出
                        break
                        pass
                    pass
                pass

            # 注意不要直接使用数据集中的变量，这样会直接修改 dataloader 中的内容
            inputs_plus = copy.deepcopy(inputs_plus)

            # 根据 se_r 的内容更新 input 的 attention_mask，把对应 aspect 的 attention 变为 0
            for aspect_start_i, aspect_end_i in se_r:
                inputs_plus['attention_mask'][0][1 + aspect_start_i:1 + aspect_end_i + 1] = 0
                pass

            # 为每个 aspect 分别计算其 opinion 和 sentiment
            for aspect_span in se_r:
                # 将原本的 inputs 的 attention_mask 进行修改，放到模型中进行训练
                as_input = torch.tensor([0] * (token_seq_len - 1)).unsqueeze(0).to(args.device)
                as_input[0][aspect_span[0]] = 1
                ae_input = torch.tensor([0] * (token_seq_len - 1)).unsqueeze(0).to(args.device)
                ae_input[0][aspect_span[1]] = 1

                inputs_plus['attention_mask'][0][1 + aspect_span[0]:1 + aspect_span[1] + 1] = 1
                # 此时相当于只关注了这一个 aspect
                y = model(inputs_plus, as_input, ae_input, args)  # 对于每一个 aspect_span 都去拿这个 span 的 opinion_sapn
                inputs_plus['attention_mask'][0][1 + aspect_span[0]:1 + aspect_span[1] + 1] = 0

                os_p, oe_p, s_logits = y['os_p'], y['oe_p'], y['s_logits']

                a_span_str = get_aspect_str_in_original_text(batch_i, aspect_span[0], aspect_span[1], token2id)
                # 这里的 a_span 是一个字符串，对应的是 batch_i - start_i - end_i
                aspect_set.add(a_span_str)

                o_spans = logits2span_list(
                    os_p.squeeze(),
                    oe_p.squeeze(),
                    3,
                    3,
                    args.b,
                    args.b_ww,
                    args.b,
                    token2id,
                    is_test
                )

                o_spans = get_opinion_spans(batch_i, o_spans, token2id)
                # o_spans 是一个 str 的 list，其中每个元素都是类似于 batch_i-start-end 的字符串
                for o_span_str in o_spans:
                    opinion_set.add(o_span_str)
                    pass

                s = torch.argmax(s_logits, dim=1).squeeze().item()
                # s 代表的是情感
                # s_logits.shape=(batch_size, 3) 当 batch_size=1 的时候，这个 s 就是一个数字，可能为 0，1，2
                as_pair, pair, triplets = get_as_pair_and_triplets(a_span_str, o_spans, s)
                # as_pair：是一个 list，其中每个元素是 i-al-ar-s，这里的 s 是数字 0、1、2 的其中一个
                # pair：是一个 list，其中每个元素是 i-al-ar-ol-or
                # triplets：是一个 list，其中每个元素是 i-al-ar-ol-or-s，这里的 s 是数字 0、1、2 的其中一个
                for token_range in pair:
                    pair_set.add(token_range)  # pair_set 是 ["i-al-ar-ol-or"]
                    pass
                for token_range in triplets:
                    triplets_set.add(token_range)  # triplets_set 是 [“i-al-ar-ol-or-s"]
                    if batch_i in multi_aspect_id:
                        multi_set.add(token_range)
                        pass
                    else:
                        single_set.add(token_range)
                        pass
                    pass

                for token_range in as_pair:
                    as_set.add(token_range)  # as_set 是 [”i-al-ar-s“]
                    pass
                pass
            pass
        pass

    a_p, a_r, a_f = score_set(aspect_set, three_goden_set[0])  # original_data.get_all_aspect_set(),
    o_p, o_r, o_f = score_set(opinion_set, three_goden_set[1])  # original_data.get_all_opinion_set(),
    t_p, t_r, t_f = score_set(triplets_set, three_goden_set[2][0])  # triplets_set 是 [“i-al-ar-ol-or-s"] 第二个是 original_data.all_triplets,
    mt_p, mt_r, mt_f = score_set(multi_set, three_goden_set[2][1])  # original_data.all_multi_triplets,
    st_p, st_r, st_f = score_set(single_set, three_goden_set[2][2])  # original_data.all_single_triplets,

    logger.info(get_dict_str_in_lines({
        "aspect.p r f1": " ".join(['%.5f' % a_p, '%.5f' % a_r, '%.5f' % a_f]),
        "opinion.p r f1": " ".join(['%.5f' % o_p, '%.5f' % o_r, '%.5f' % o_f]),
        "multi_t.p r f1": " ".join(['%.5f' % mt_p, '%.5f' % mt_r, '%.5f' % mt_f]),
        "single_t.p r f1": " ".join(['%.5f' % st_p, '%.5f' % st_r, '%.5f' % st_f]),
        "all_t.p r f1": " ".join(['%.5f' % t_p, '%.5f' % t_r, '%.5f' % t_f]),
    }))

    return mt_f, st_f, t_f
